"""
Calculate pre-miRNA expression based on high throughput sequencing data.

Usage:
  premirna_expression <premirna_reference> <reads> ...

Options:
  -h --help      Show this screen.
"""

from __future__ import print_function
import sys, subprocess, docopt, re, datetime, pwd, os
from collections import Counter
from pypette import shell, read_fasta

def premirna_expression(read_paths, premirna_reference_path):
	
	colorspace = True
	
	tmp_path = ''
	if not tmp_path:
		tmp_path = '/data/tmp/%s/premirna_expression_%s' % (
			pwd.getpwuid(os.getuid()).pw_name,
			datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S'))
		os.makedirs(tmp_path)
	
	premirna = read_fasta(premirna_reference_path)
	sorted_premirna = sorted(premirna.iterkeys())
	
	premirna_index = tmp_path + '/premirnas'
	
	shell('bowtie-build -q %s %s %s' % 
		(' -C' if colorspace else '', premirna_reference_path, premirna_index))
	
	# Write out the pre-miRNA names and sequences
	report = open(tmp_path + '/premirnas.tsv', 'w')
	report.write('NAME\tSEQUENCE\n')
	for key in sorted_premirna:
		report.write('%s\t%s\n' % (key, premirna[key]))
	report.close()
	
	for s, read_path in enumerate(read_paths):
		sample_name = os.path.basename(read_path)
		alignments_path = '%s/%s.bowtie' % (tmp_path, sample_name)
		
		if read_path.lower().endswith('.gz'):
			read_path = '<(gunzip -c %s)' % read_path
		
		shell('bowtie -C -f -p8 --best --norc --trim3 5 -v2 -k1 -m1 %s %s '
			'--suppress=1,2,7,8 > %s' % 
			(premirna_index, read_path, alignments_path))
		
		premirna_reads = Counter({key: 0 for key in premirna})
		alignments = open(alignments_path)
		for line in alignments:
			tokens = line[:-1].split('\t')
			premirna_reads[tokens[0]] += 1
	
		# Write out the expression levels in this sample
		report = open('%s/%s_expr.tsv' % (tmp_path, sample_name), 'w')
		report.write('%s\n' % sample_name)
		for key in sorted_premirna:
			report.write('%d\n' % premirna_reads[key])
		report.close()
	
	
	
		
	
	
if __name__ == '__main__':
	args = docopt.docopt(__doc__)
	
	premirna_expression(args['<reads>'], args['<premirna_reference>'])
	

